/* Copyright Thanh-Trung Pham, JAIST, 2013 */

package jp.ac.jaist.model;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.math.linear.Array2DRowRealMatrix;
import org.apache.commons.math.linear.LUDecompositionImpl;
import org.apache.commons.math.linear.MatrixUtils;
import org.apache.commons.math.linear.RealMatrix;

/**
 *
 * @author Thanh-Trung Pham
 */
public class RMPMarkovModel {

    private RMPMarkovState m_oSTARTState;
    private Map<String, RMPMarkovState> m_oNameTransientStateMappings;
    private Map<String, RMPMarkovState> m_oNameAbsorbingStateMappings;
    private List<RMPMarkovTransition> m_oTransitions;
    private double[][] m_oMarkovChain;
    private double[][] m_oBMatrix;

    public RMPMarkovModel() {
        m_oNameTransientStateMappings = new HashMap<String, RMPMarkovState>();
        m_oNameAbsorbingStateMappings = new HashMap<String, RMPMarkovState>();
        m_oTransitions = new ArrayList<RMPMarkovTransition>();
    }

    /**
     * Create a START state.
     * @param startStateName the START state name.
     * @return true if created successfully, otherwise false.
     */
    public boolean createSTARTState(String startStateName) {
        if (m_oSTARTState != null) {
            return false;
        }

        if (StringUtils.isBlank(startStateName)) {
            return false;
        }

        m_oSTARTState = new RMPMarkovState(startStateName);

        return true;
    }

    /**
     * Check if this Markov model contains a state with the given name.
     * @param stateName the given state name.
     * @return true if this Markov model contains a state with the given name, otherwise false.
     */
    private boolean containState(String stateName) {
        if (m_oSTARTState != null) {
            if (StringUtils.equals(m_oSTARTState.getStateName(), stateName)) {
                return true;
            }
        }

        if (m_oNameTransientStateMappings.containsKey(stateName) || m_oNameAbsorbingStateMappings.containsKey(stateName)) {
            return true;
        }

        return false;
    }

    /**
     * Get a Markov state by name.
     * @param stateName a state name.
     * @return the Markov state with the given name.
     */
    private RMPMarkovState getStateByName(String stateName) {
        if (StringUtils.equals(m_oSTARTState.getStateName(), stateName)) {
            return m_oSTARTState;
        }

        if (m_oNameTransientStateMappings.containsKey(stateName)) {
            return m_oNameTransientStateMappings.get(stateName);
        }

        return m_oNameAbsorbingStateMappings.get(stateName);
    }

    /**
     * Add a new transient Markov state.
     * @param stateName the state name.
     * @return true if added successfully, otherwise false.
     */
    public boolean addTransitiveMarkovState(String stateName) {
        if (StringUtils.isBlank(stateName)) {
            return false;
        }

        if (containState(stateName)) {
            return false;
        }

        RMPMarkovState oState = new RMPMarkovState(stateName);
        m_oNameTransientStateMappings.put(stateName, oState);

        return true;
    }

    /**
     * Add a new absorbing Markov state.
     * @param stateName the state name.
     * @return true if added successfully, otherwise false.
     */
    public boolean addAbsorbingMarkovState(String stateName) {
        if (StringUtils.isBlank(stateName)) {
            return false;
        }

        if (containState(stateName)) {
            return false;
        }

        RMPMarkovState oState = new RMPMarkovState(stateName);
        m_oNameAbsorbingStateMappings.put(stateName, oState);

        return true;
    }

    /**
     * Add a new Markov transition.
     * @param sourceStateName the source state name.
     * @param destinationStateName the destination state name.
     * @param transitionProbability the transition probability.
     * @return true if added successfully, otherwise false.
     */
    public boolean addMarkovTransition(String sourceStateName, String destinationStateName, double transitionProbability) {
        if (transitionProbability < 0.0 || transitionProbability > 1.0) {
            return false;
        }

        if (StringUtils.isBlank(sourceStateName) || StringUtils.isBlank(destinationStateName)) {
            return false;
        }

        if (!containState(sourceStateName) || !containState(destinationStateName)) {
            return false;
        }

        RMPMarkovState oSourceState = getStateByName(sourceStateName);
        RMPMarkovState oDestinationState = getStateByName(destinationStateName);

        RMPMarkovTransition oTransition = new RMPMarkovTransition(oSourceState, oDestinationState, transitionProbability);

        if (m_oTransitions.contains(oTransition)) {
            return false;
        }

        m_oTransitions.add(oTransition);

        return true;
    }

    /**
     * Generate the Markov chain.
     */
    private void generateMarkovChain() {
        int iStateCount = 0;

        m_oSTARTState.setID(iStateCount++);

        Iterator<RMPMarkovState> oIterator;

        oIterator = m_oNameTransientStateMappings.values().iterator();
        while (oIterator.hasNext()) {
            RMPMarkovState oState = oIterator.next();
            oState.setID(iStateCount++);
        }

        int iAbsorbingStateCount = 0;
        oIterator = m_oNameAbsorbingStateMappings.values().iterator();
        while (oIterator.hasNext()) {
            RMPMarkovState oState = oIterator.next();
            oState.setID(iStateCount++);
            oState.setAbsorbingID(iAbsorbingStateCount++);
        }

        m_oMarkovChain = new double[iStateCount][iStateCount];
        for (int i = 0; i < iStateCount; i++) {
            for (int j = 0; j < iStateCount; j++) {
                m_oMarkovChain[i][j] = 0.0;
            }
        }

        for (int i = 0; i < m_oTransitions.size(); i++) {
            RMPMarkovTransition oTransition = m_oTransitions.get(i);
            m_oMarkovChain[oTransition.getSourceState().getID()][oTransition.getDestinationState().getID()] = oTransition.getTransitionProbability();
        }
    }

    /**
     * Solve the Markov chain.
     */
    private void solveMarkovChain() {
        RealMatrix oMarkovChain = new Array2DRowRealMatrix(m_oMarkovChain);

        int iNumberOfENDState = m_oNameAbsorbingStateMappings.size();

        RealMatrix oQMatrix = oMarkovChain.getSubMatrix(0, oMarkovChain.getRowDimension() - (iNumberOfENDState + 1), 0, oMarkovChain.getColumnDimension() - (iNumberOfENDState + 1));
        RealMatrix oRMatrix = oMarkovChain.getSubMatrix(0, oMarkovChain.getRowDimension() - (iNumberOfENDState + 1), oMarkovChain.getColumnDimension() - iNumberOfENDState, oMarkovChain.getColumnDimension() - 1);

        RealMatrix oIdentityMatrix = MatrixUtils.createRealIdentityMatrix(oQMatrix.getRowDimension());

        m_oBMatrix = ((new LUDecompositionImpl(oIdentityMatrix.subtract(oQMatrix)).getSolver().getInverse()).multiply(oRMatrix)).getData();
    }

    /**
     * Generate and then solve the Markov chain.
     */
    private void generateAndSolveMarkovChain() {
        if (m_oBMatrix == null) {
            generateMarkovChain();
            solveMarkovChain();
        }
    }

    /**
     * Get the probability of reaching a absorbing state with the given name from the START state.
     * @param absorbingStateName the given name of the absorbing state.
     * @return the probability of reaching a absorbing state with the given name from the START state.
     */
    public Double getProbability(String absorbingStateName) {
        generateAndSolveMarkovChain();

        if (!m_oNameAbsorbingStateMappings.containsKey(absorbingStateName)) {
            return null;
        }

        return m_oBMatrix[0][getStateByName(absorbingStateName).getAbsorbingID()];
    }
}
